package com.gaga.console.base.utils;

import com.alibaba.fastjson.JSON;
import com.gaga.common.enums.PushMessageType;
import com.gaga.console.base.entity.sse.PushMessageDto;
import com.gaga.metadata.common.utils.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

/**
 * @author liujia
 */
@Slf4j
@Component
public class SseEmitterServerUtils {

    private static final Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();
    
    // 用于跟踪连接健康状态
    private static final Map<String, AtomicInteger> connectionHealthMap = new ConcurrentHashMap<>();
    
    // 最大心跳失败次数，超过此次数认为连接已断开
    private static final int MAX_HEARTBEAT_FAILURES = 3;
    
    /**
     * 心跳消息内容
     */
    private static final PushMessageDto HEARTBEAT_MESSAGE = new PushMessageDto(PushMessageType.HEARTBEAT, null, null, null, null, 0);

    /**
     * 心跳间隔时间，单位毫秒（15秒）
     */
    private static final long HEARTBEAT_INTERVAL = 15_000L;

    /**
     * 创建新的SSE连接
     */
    public static SseEmitter connect(String channelId) {
        //设置超时时间，0表示不过期，默认是30秒，超过时间未完成会抛出异常
        SseEmitter sseemitter = new SseEmitter(0L);
        
        //注册回调
        sseemitter.onCompletion(completionCallBack(channelId));
        sseemitter.onError(errorCallBack(channelId));
        sseemitter.onTimeout(timeoutCallBack(channelId));
        
        sseEmitterMap.put(channelId, sseemitter);
        connectionHealthMap.put(channelId, new AtomicInteger(0));
        
        // 连接建立后立即发送一次心跳
        try {
            sseemitter.send(HEARTBEAT_MESSAGE);
            log.info("create new sse connect, sent initial heartbeat, channelId: {}", channelId);
        } catch (IOException e) {
            log.error("Failed to send initial heartbeat to channelId: {}, error: {}", channelId, e.getMessage());
            // 不要在此处移除连接，让错误回调处理
        }
        
        return sseemitter;
    }

    /**
     * 关闭指定的SSE连接
     */
    public static void close(String channelId) {
        SseEmitter sseEmitter = getId(channelId);
        if (sseEmitter != null) {
            try {
                sseEmitter.complete();
            } catch (Exception e) {
                log.error("Error closing SSE connection for channelId: {}, error: {}", channelId, e.getMessage());
            } finally {
                remove(channelId);
            }
        }
    }

    /**
     * 批量发送消息给所有连接
     */
    public static void batchSendMessage(Object message) {
        Set<Map.Entry<String, SseEmitter>> entries = sseEmitterMap.entrySet();
        for (Map.Entry<String, SseEmitter> entry : entries) {
            String channelId = entry.getKey();
            SseEmitter emitter = entry.getValue();
            try {
                emitter.send(message);
                // 成功发送心跳，重置失败计数
                if (message == HEARTBEAT_MESSAGE) {
                    resetHeartbeatFailures(channelId);
                }
            } catch (IOException e) {
                log.error("channelId: {}, send message error: {}", channelId, e.getMessage());
                
                // 如果是心跳消息，增加失败计数
                if (message == HEARTBEAT_MESSAGE) {
                    incrementHeartbeatFailures(channelId);
                } else {
                    // 如果是普通消息发送失败，移除连接
                    remove(channelId);
                }
            }
        }
    }

    /**
     * 给指定用户发消息
     */
    public static void sendMessage(String channelId, Object message) {
        if (StringUtils.isBlank(channelId)) {
            return;
        }
        
        if (sseEmitterMap.containsKey(channelId)) {
            try {
                sseEmitterMap.get(channelId).send(message);
                // 成功发送，如果是心跳，重置计数
                if (message == HEARTBEAT_MESSAGE) {
                    resetHeartbeatFailures(channelId);
                }
            } catch (Exception e) {
                log.error("channelId: {}, send message error: {}, message: {}", 
                          channelId, e.getMessage(), JSON.toJSONString(message));
                
                // 如果是心跳消息，增加失败计数
                if (message == HEARTBEAT_MESSAGE) {
                    incrementHeartbeatFailures(channelId);
                } else {
                    // 非心跳消息发送失败直接移除
                    remove(channelId);
                }
            }
        }
    }

    /**
     * 移除用户连接
     */
    public static SseEmitter remove(String channelId) {
        if (sseEmitterMap.containsKey(channelId)) {
            SseEmitter removed = sseEmitterMap.remove(channelId);
            connectionHealthMap.remove(channelId);
            log.info("removed channelId: {}", channelId);
            return removed;
        }
        return null;
    }

    /**
     * 定时发送心跳消息
     */
    @Scheduled(fixedRate = HEARTBEAT_INTERVAL)
    public void sendHeartbeat() {
        if (sseEmitterMap.isEmpty()) {
            log.debug("No active SSE connections, skipping heartbeat");
            return;
        }
        
        log.debug("Sending heartbeat to all connected clients ({})", sseEmitterMap.size());
        batchSendMessage(HEARTBEAT_MESSAGE);
        
        // 检查并清理失败的连接
        cleanupFailedConnections();
    }
    
    /**
     * 检查并清理心跳失败的连接
     */
    private void cleanupFailedConnections() {
        Set<Map.Entry<String, AtomicInteger>> entries = connectionHealthMap.entrySet();
        for (Map.Entry<String, AtomicInteger> entry : entries) {
            String channelId = entry.getKey();
            AtomicInteger failures = entry.getValue();
            
            if (failures.get() >= MAX_HEARTBEAT_FAILURES) {
                log.warn("Channel {} exceeded max heartbeat failures ({}), closing connection", 
                         channelId, failures.get());
                close(channelId);
            }
        }
    }
    
    /**
     * 增加心跳失败计数
     */
    private static void incrementHeartbeatFailures(String channelId) {
        AtomicInteger counter = connectionHealthMap.get(channelId);
        if (counter != null) {
            int newCount = counter.incrementAndGet();
            log.warn("Heartbeat failure for channelId: {}, failure count: {}/{}", 
                     channelId, newCount, MAX_HEARTBEAT_FAILURES);
        }
    }
    
    /**
     * 重置心跳失败计数
     */
    private static void resetHeartbeatFailures(String channelId) {
        AtomicInteger counter = connectionHealthMap.get(channelId);
        if (counter != null && counter.get() > 0) {
            counter.set(0);
            log.debug("Reset heartbeat failures for channelId: {}", channelId);
        }
    }

    /**
     * 获取指定连接
     */
    public static SseEmitter getId(String channelId) {
        return sseEmitterMap.get(channelId);
    }
    
    /**
     * 获取当前活跃连接数
     */
    public static int getActiveConnectionCount() {
        return sseEmitterMap.size();
    }

    /**
     * 连接完成回调
     */
    private static Runnable completionCallBack(String channelId) {
        return () -> {
            log.info("SSE connection completed: {}", channelId);
            remove(channelId);
        };
    }

    /**
     * 连接超时回调
     */
    private static Runnable timeoutCallBack(String channelId) {
        return () -> {
            log.info("SSE connection timeout: {}", channelId);
            remove(channelId);
        };
    }

    /**
     * 连接错误回调
     */
    private static Consumer<Throwable> errorCallBack(String channelId) {
        return (throwable) -> {
            log.error("SSE connection error: {}, error: {}", channelId, throwable.getMessage());
            remove(channelId);
        };
    }
}

